package com.sicau.agent;

import cn.hutool.core.util.StrUtil;
import com.sicau.agent.model.AgentState;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 处理工具调用的具体代理类，实现了ReActAgent的think·act方法`
 */
@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent {
    /**
     * 聊天选项, 手动控制工具执行流程
     */
    private ChatOptions chatOptions;
    /**
     * 可用工具集合
     */
    private ToolCallback[] toolCallbacks;
    /**
     * 工具调用管理器
     */
    private ToolCallingManager toolCallingManager;


    public ToolCallAgent(ToolCallback[] toolCallbacks) {
        this.toolCallingManager = ToolCallingManager.builder().build();
        this.chatOptions = ToolCallingChatOptions.builder()
                .toolCallbacks(toolCallbacks)
                .internalToolExecutionEnabled(false)
                .build();
    }

    @Override
    protected boolean think() {
        List<Message> messageList = getMessageList();
        if (!StrUtil.isBlank(getNextStepPrompt())) {
            messageList.add(new UserMessage(getNextStepPrompt()));
        }
        try {
            Prompt prompt = new Prompt(messageList, chatOptions);
            ChatClient chatClient = getChatClient();
            // 获取带工具选项的响应
            ChatResponse chatResponse = chatClient.prompt(prompt)
                    .system(getSystemPrompt())
                    .call()
                    .chatResponse();
            setChatResponse(chatResponse);
            // 输出提示信息
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            log.info("{}的思考：{}", getName(),  assistantMessage.getText());
            log.info("{}一共调用了{}个工具", getName(), toolCallList.size());
            String toolsInfo = toolCallList
                    .stream()
                    .map(toolCall ->
                            String.format("工具名称：%s, 工具参数：%s", toolCall.name(), toolCall.arguments()))
                    .collect(Collectors.joining("\n"));
            log.info(toolsInfo);
            if (toolCallList.isEmpty()) {
                // 只有不调用工具时才记录消息
                getMessageList().add(assistantMessage);
                return false;
            } else {
                return true;
            }
        } catch (Exception e) {
            log.error("{}思考过程出现了一点问题！", getName(), e);
            getMessageList().add(new AssistantMessage("思考过程中出现了一些问题！"));
            return false;
        }
    }

    @Override
    protected String act() {
        try {
            ChatResponse chatResponse = getChatResponse();
            // 调用工具
            ToolExecutionResult toolExecutionResult = toolCallingManager
                    .executeToolCalls(new Prompt(getMessageList(), chatOptions),
                            chatResponse);
            // 记录消息上下文
            List<Message> messages = toolExecutionResult.conversationHistory();
            setMessageList(messages);
            // 当前工具调用结果
            ToolResponseMessage toolResponseMessage = (ToolResponseMessage) messages.getLast();
            String toolExecutionResults = toolResponseMessage
                    .getResponses()
                    .stream()
                    .map(toolResponse ->
                            String.format("工具%s被调用了，调用结果: \n%s", toolResponse.name(), toolResponse.responseData()))
                    .collect(Collectors.joining("\n"));
            // 判断是否调用了终止工具
            boolean isTerminate = toolResponseMessage
                    .getResponses()
                    .stream()
                    .anyMatch(toolResponse -> toolResponse.name().equals("doTerminate"));
            if (isTerminate) {
                setAgentState(AgentState.FINISHED);
            }
            log.info(toolExecutionResults);
            return toolExecutionResults;
        } catch (Exception e) {
            log.error("{}执行工具过程出现一些问题！", getName(), e);
            return "执行工具过程出现一些问题！";
        }
    }
}
